#!/usr/bin/env python
import re

''' This module provide the class to parse junction, sam and wig file
generated from topHat '''
class JunctionLineInfo( object ):
    def __init__( self, line ):
        line = line.strip()
        info = line.split( "\t" )
        MIN_RECORD_USED = 5
        if len( info ) >= MIN_RECORD_USED:
            try:
                self.gene = info[0]
                self.start = int( info[1] )
                self.end = int( info[2] )
                #self.fakeName = extn[3] 
                self.noMaps = info[4]
                self.strand = info[5]
            except ValueError:
                print self.gene
                print "has wrong value"
        else:
            pass;
    def out_line( self, shift ):
        '''Output the bed format line, with position shift'''
        line = self.gene + '\t'\
            + str( self.start - shift ) + '\t'\
            + str( self.end - shift ) + '\t'\
            + self.noMaps + '\t'\
            + self.strand + '\n'
        return line

class SamLineInfo( object ):
    '''
    Object contains 
    QNAME, FLAG, RNAME, POS, MAPQ, CIGAR, (SEQ) ...
    __init__ takes a line as input
    '''
    def __init__( self, line ):
        line = line.strip()
        extn = line.split( "\t" )
        col_num = len( extn )
        if col_num > 5:
            try:
                self.QNAME = extn[0]
                self.FLAG = extn[1]
                # Reference sequence Name
                self.RNAME = extn[2]
                self.POS = int ( extn[3] )
                self.MAPQ = int ( extn[4] )
                self.CIGAR = extn[5]
                if( col_num >= 10 ):
                    self.MRNM = extn[6]
                    self.MPOS = int( extn[7] )
                    self.ISIZE = int( extn[8] )
                    self.SEQ = extn[9]
                else:
                    self.SEQ = ""
            except ValueError:
                    print self.QNAME
                    print "has wrong value"
        else:
            pass;

    def out_line( self, shift ):
        '''Output the sam format line, with position shift'''
        line = self.QNAME + '\t'\
            + self.FLAG + '\t'\
            + self.RNAME + '\t'\
            + str( self.POS - shift ) + '\t'\
            + str( self.MAPQ ) + '\t'\
            + self.CIGAR
        # This checking is to save the space, which should be removed later 
        if ( self.SEQ != "" ):
            line = line + '\t'\
            + self.MRNM + '\t'\
            + str( self.MPOS ) + '\t'\
            + str( self.ISIZE ) + '\t'\
            + self.SEQ
        line += '\n'
        return line

    def get_junction_from_CIGAR( self ):
        ''' If the alignment supports a junction, 
            extract the junction start, end and strand info.
            (For non-mate pair only) 
        '''
        cigarP = re.compile( '[0-9]*M[0-9]*N[0-9]*M' )
        try:
            if ( cigarP.match( self.CIGAR ) ):
                self.CIGAR = self.CIGAR.replace( 'N', '\t' )
                self.CIGAR = self.CIGAR.replace( 'M', '\t' )
                self.CIGAR = self.CIGAR.split( '\t' )
                # double check the start and end
                self.junction_start = self.POS + int( self.CIGAR[0] )
                self.junction_end = self.junction_start + int ( self.CIGAR[1] )
                if( int( self.FLAG ) & 0x10 == 0 ):
                    self.strand = '+'
                else:
                    self.strand = '-'
                return True
            else:
                return False # not junction
        except:
            return False # not junction or with format error

class WigLineInfo( object ):
    def __init__( self, line ):
        line = line.strip()
        extn = line.split( "\t" )
        if len( extn ) >= 3:
            try:
                self.gene = extn[0]
                self.start = int( extn[1] )
                self.end = int( extn[2] )
                self.coverage = int ( extn[3] )
            except ValueError:
                print self.gene
                print "has wrong value"
        else:
            pass;

    def out_line( self, shift ):
        '''Output the wig format line, with position shift'''
        line = self.gene + '\t'\
            + str( self.start - shift ) + '\t'\
            + str( self.end - shift ) + '\t'\
            + str( self.coverage ) + '\n'
        return line
